
import pandas as pd
import numpy as np
import csv
from tqdm import tqdm

# 配置参数
TOTAL_RECORDS = 10_000_000
BATCH_SIZE = 100_000
START_USER_ID = 1_000_001
FILE_NAME = "user_profile_data.csv"

# 城市分布配置
CITY_WEIGHTS = {
    # 一线城市 30%
    "北京": 0.08, "上海": 0.08, "广州": 0.07, "深圳": 0.07,
    # 二线城市 50%
    "杭州": 0.06, "成都": 0.06, "武汉": 0.05, "南京": 0.05,
    "苏州": 0.05, "重庆": 0.05, "天津": 0.04, "西安": 0.04,
    "郑州": 0.03, "长沙": 0.03, "沈阳": 0.02, "青岛": 0.02,
    # 三线城市 20%
    "昆明": 0.02, "福州": 0.02, "济南": 0.02, "合肥": 0.02,
    "宁波": 0.02, "南宁": 0.02, "哈尔滨": 0.02, "石家庄": 0.02,
    "贵阳": 0.02, "太原": 0.02
}

# 手机品牌分布
PHONE_BRANDS = ["华为", "苹果", "小米", "OPPO", "vivo", "三星"]
BRAND_PROBS = [0.28, 0.22, 0.20, 0.15, 0.10, 0.05]


def generate_batch(start_id, batch_size):
    """生成单批次数据"""
    user_ids = np.arange(start_id, start_id + batch_size)

    # 年龄生成（正态分布）
    ages = np.random.normal(loc=35, scale=10, size=batch_size)
    ages = np.clip(np.round(ages), 18, 70).astype(int)

    # 性别生成（含1%空值）
    genders = np.random.choice(
        ["男", "女", np.nan],
        size=batch_size,
        p=[0.52, 0.47, 0.01]
    )

    # 城市生成（加权随机）
    cities = np.random.choice(
        list(CITY_WEIGHTS.keys()),
        size=batch_size,
        p=list(CITY_WEIGHTS.values())
    )

    # 手机品牌
    phone_brands = np.random.choice(
        PHONE_BRANDS,
        size=batch_size,
        p=BRAND_PROBS
    )

    # 浏览时长（对数正态分布 + 长尾）
    base_view_time = np.random.lognormal(mean=8.5, sigma=1.2, size=batch_size)
    long_tail = np.random.pareto(0.5, size=batch_size) * 3600 * 2
    view_times = np.clip((base_view_time * 1800 + long_tail), 0, 180 * 86400).astype(int)

    # 浏览次数（泊松分布 + 长尾）
    view_counts = np.random.poisson(lam=150, size=batch_size)
    view_counts += np.random.pareto(0.7, size=batch_size).astype(int)
    view_counts = np.clip(view_counts, 0, 50000)

    # 交易次数（零膨胀泊松分布）
    has_order = np.random.random(size=batch_size) > 0.65
    order_counts = np.where(
        has_order,
        np.random.poisson(lam=3.5, size=batch_size) + 1,
        0
    )
    order_counts = np.clip(order_counts, 0, 100)

    # 交易金额（基于交易次数生成）
    base_amount = np.random.normal(300, 50, size=batch_size)
    order_amounts = np.round(
        np.clip(base_amount, 100, 1000) * order_counts,
        2
    )

    # 最近交易天数（有交易用户生成）
    recent_days = np.where(
        order_counts > 0,
        np.random.randint(1, 181, size=batch_size),
        None
    )

    # 组装批次数据
    return zip(
        user_ids,
        ages,
        genders,
        cities,
        phone_brands,
        view_times,
        view_counts,
        order_amounts,
        order_counts,
        recent_days
    )


# 生成CSV文件
with open(FILE_NAME, 'w', newline='', encoding='utf-8') as csvfile:
    writer = csv.writer(csvfile)
    # 写入表头
    writer.writerow([
        'user_id', 'age', 'gender', 'region_city', 'phone_brand',
        'page_view_during_time_180d', 'page_view_count_180d',
        'order_amount_180d', 'order_count_180d', 'order_recent_days_180d'
    ])

    # 分批生成数据
    for i in tqdm(range(0, TOTAL_RECORDS, BATCH_SIZE)):
        batch = generate_batch(START_USER_ID + i, min(BATCH_SIZE, TOTAL_RECORDS - i))
        writer.writerows(batch)

print(f"数据生成完成！文件路径: {FILE_NAME}")